#! /usr/bin/env python
'''
Pour water from mug to bowl.
'''
import sys, os, time, yaml, copy
src_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../..')
sys.path.append(src_dir)
import rospy
import ros_np_multiarray as ros_nm
import numpy as np
from geometry_msgs.msg import PoseStamped
from visualization_msgs.msg import Marker, MarkerArray
import moveit_commander
import tf2_ros
############# perception module #################
sys.path.insert(0, '{}/perception/vision_perception/src/'.format(src_dir))
# from perception.vision_perception.src.perception_client import Perception_Client
from perception_client import Perception_Client
sys.path.pop(0)
sys.path.insert(0, '{}/perception/vision_grasp_est/src/'.format(src_dir))
# from perception.vision_grasp_est.src.grasp_est_client import GraspEstClient
from grasp_est_client import GraspEstClient
sys.path.pop(0)
############# planning module ###################
from planning.scripts.action_client import *
#################################################

from utilities.tools.vis_utils import *
from utilities.tools.tf_utils import R2q, get_pose_list_from_pose_stamped, rot_Y, transform_poselist_to_tgt


def get_pour_pos_for_grasp_handle(mug_info, bowl_info):
    '''
    args:
        mug_info: SingleObjPose
        bowl_info: SingleObjPose
    return:
        pour_pos: list, [position, ]
    '''
    mug_box = mug_info.box
    bowl_box = bowl_info.box
    mug_center_x = mug_box[0] + mug_box[2]
    bowl_center_x = bowl_box[0] + bowl_box[2]
    mug_pose = mug_info.RT
    mug_pose = ros_nm.to_numpy_f32(mug_pose)
    bowl_pose = bowl_info.RT
    bowl_pose = ros_nm.to_numpy_f32(bowl_pose)
    mug_radius = mug_info.scales[2] * mug_info.scale / 2
    mug_radius_handle = mug_info.scales[0] * mug_info.scale / 2
    bowl_radius = bowl_info.scales[2] * bowl_info.scale / 2
    bowl_height = bowl_info.scales[1] * bowl_info.scale / 2
    mug_pose_z = mug_pose[:,2]
    mug_pose_x = mug_pose[:,0]
    bowl_pose_y = bowl_pose[:,1]
    pour_pos = copy.deepcopy(bowl_pose[:,3])
    # supposed that gripper is right hand, so it pours on the right of bowl
    if mug_center_x > bowl_center_x:
        right = True
        pour_pos += -(bowl_radius + mug_radius - 0.02) * mug_pose_z
    else:
        right = False
        pour_pos -= -(bowl_radius + mug_radius - 0.02) * mug_pose_z
    pour_pos += (bowl_height + 0.14) * bowl_pose_y
    pour_pos += (mug_radius_handle * 0.1) * mug_pose_x

    return pour_pos.tolist(), right

def get_pour_pose_for_grasp_body(mug_info, bowl_info):
    '''
    args:
        mug_info: SingleObjPose
        bowl_info: SingleObjPose
    return:
        pour_pos: list, [position, ]
    '''
    mug_box = mug_info.box
    bowl_box = bowl_info.box
    mug_center_x = mug_box[0] + mug_box[2]
    bowl_center_x = bowl_box[0] + bowl_box[2]
    mug_pose = mug_info.RT
    mug_pose = ros_nm.to_numpy_f32(mug_pose)
    bowl_pose = bowl_info.RT
    bowl_pose = ros_nm.to_numpy_f32(bowl_pose)
    mug_radius = mug_info.scales[2] * mug_info.scale / 2
    mug_radius_handle = mug_info.scales[0] * mug_info.scale / 2
    bowl_radius = bowl_info.scales[2] * bowl_info.scale / 2
    bowl_height = bowl_info.scales[1] * bowl_info.scale / 2
    mug_pose_z = mug_pose[:,2]
    mug_pose_x = mug_pose[:,0]
    bowl_pose_x = bowl_pose[:,0]
    bowl_pose_y = bowl_pose[:,1]
    bowl_pose_z = bowl_pose[:,2]
    pour_pos = copy.deepcopy(bowl_pose[:,3])
    # supposed that gripper is right hand, so it pours on the right of bowl
    # pour_pos += -(bowl_radius + mug_radius) * mug_pose_z
    
    if mug_center_x > bowl_center_x:
        right = True
        # pour_pos += (bowl_radius + mug_radius - 0.02) * bowl_pose_x
        pour_rot = np.matmul(bowl_pose[:3,:3], rot_Y(np.pi * 7 / 6))
        pour_pos -= (0.03) * bowl_pose_z
        pour_pos += (0.05) * bowl_pose_x
    else:
        right = False
        # pour_pos -= (bowl_radius + mug_radius - 0.02) * bowl_pose_x
        pour_rot = np.matmul(bowl_pose[:3,:3], rot_Y(np.pi * 5 / 6))
        pour_pos -= (0.03) * bowl_pose_z
        pour_pos -= (0.05) * bowl_pose_x
    pour_qua = R2q(pour_rot)
    pour_pos += (bowl_height + 0.13) * bowl_pose_y

    return pour_pos.tolist(), pour_qua.tolist(), right

if __name__ == '__main__':

    vis = True
    grasp_flag = True
    robot_prefix = 'j2n6s300_'
    camera_frame = 'camera_color_optical_frame'
    src_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../..')
    rospy.init_node('pour_water_node')
    tfBuffer = tf2_ros.Buffer()
    tf_listenser = tf2_ros.TransformListener(tfBuffer)
    arm_group = moveit_commander.MoveGroupCommander("arm")
    # homeRobot(robot_prefix)
    initial_joint_pose(robot_prefix)
    grasp_est_client = GraspEstClient()

    with open(os.path.join(src_path, 'perception/vision_grasp_est/src/panda_grasp_data/pour_water_grasp.yaml'), 'r') as stream:
        grasp_cfg = yaml.safe_load(stream)

    # vis topic
    obj_pose_topic = rospy.Publisher('/pour_water/multi_obj_pose', PoseArray, queue_size=1)
    obj_topic = rospy.Publisher('/pour_water/multi_obj_marker', MarkerArray, queue_size=1)
    gripper_grasp_pose_topic = rospy.Publisher('/pour_water/grasp_pose_marker', Marker, queue_size=1)
    gripper_pour_pose_topic = rospy.Publisher('/pour_water/pour_pose_marker', Marker, queue_size=1)

    while True:
        start_time = time.time()
        perception_client = Perception_Client()
        perception_results = perception_client.client() # msg type: SingleObjPose[]
        rospy.loginfo('2D+6D use time : {}'.format(time.time() - start_time))
        if perception_results is not None:
            
            bowl_idx = -1
            mug_idx = -1
            
            # get index of mug and bowl instances
            for i in range(len(perception_results.obj_pose_list)):
                perception_result = perception_results.obj_pose_list[i]
                obj_cate = perception_result.category_name
                if obj_cate == 'mug':
                    mug_idx = i
                elif obj_cate == 'bowl':
                    bowl_idx = i
                else:
                    continue
            
            if mug_idx == -1 or bowl_idx == -1:
                rospy.loginfo('The mug or bowl has been not detected!')
                continue
            
            # vis in rviz
            if vis:
                obj_pose_list = []
                marker_list = []
                for i in [mug_idx, bowl_idx]:
                    perception_result = perception_results.obj_pose_list[i] # msg type: SingObjPose
                    obj_pose = perception_result.RT
                    obj_pose = ros_nm.to_numpy_f32(obj_pose)
                    obj_cate = perception_result.category_name
                    rospy.loginfo('Detect Obj Category:  {}'.format(obj_cate))
                    
                    obj_position = obj_pose[:3,3].tolist()      
                    obj_orientation = R2q(obj_pose[:3,:3]).tolist()
                    obj_pose_list.append(get_pose_msg(obj_position, obj_orientation))
                    mesh_resource = 'file://{}/perception/vision_grasp_est/src/lib/shape_encoder/recon_meshes/{}_mesh.ply'.format(src_path, obj_cate)
                    scale = perception_result.scale * 0.5
                    marker_list.append(get_marker_msg(obj_position, obj_orientation, scale, i, mesh_resource, camera_frame, color=[0,0,1,1]))
                
                obj_pose_array_msg = get_posearray_msg(obj_pose_list, camera_frame)
                obj_pose_topic.publish(obj_pose_array_msg)
                obj_maker_array_msg = get_markerarray_msg(marker_list)
                obj_topic.publish(obj_maker_array_msg)

            
            grasp_obj_cate = 'mug'
            # estimate grasp pose for mug
            mug_result = perception_results.obj_pose_list[mug_idx]
            grasp_est_pose = grasp_est_client.get_grasp_pose(mug_result) # msg type: std_msgs/MultiArray and SingleObjPose[]

            if grasp_est_pose is not None:
                
                grasp_pose = ros_nm.to_numpy_f32(grasp_est_pose.grasp_pose) # grasp pose in OBJ COORDINATES [4x4]
                obj_pose = grasp_est_pose.single_obj_pose.RT # may add icp
                obj_pose = ros_nm.to_numpy_f32(obj_pose)
                grasp_pose_CAM = np.matmul(obj_pose, grasp_pose) # grasp pose in CAMERA COORDINATES, [4x4]
                
                rospy.loginfo('Total use time : {}'.format(time.time() - start_time))

                pre_position = grasp_pose_CAM[:3,3] - 0.10 * grasp_pose_CAM[:3, 2]
                pre_position = pre_position.tolist()
                depth_offset = grasp_cfg[grasp_obj_cate]['depth_offset'][grasp_cfg[grasp_obj_cate]['type']]
                position = grasp_pose_CAM[:3,3] + depth_offset * grasp_pose_CAM[:3, 2]
                position = position.tolist()
                orientation = R2q(grasp_pose_CAM[:3,:3]).tolist()

                # estimate pour pose according to bowl's pose
                bowl_result = perception_results.obj_pose_list[bowl_idx]
                if grasp_cfg['mug']['type'] == 'body':
                    pour_position, pour_orientation, right = get_pour_pose_for_grasp_body(mug_result, bowl_result)
                else:
                    pour_position, right = get_pour_pos_for_grasp_handle(mug_result, bowl_result)
                    pour_orientation = orientation

                pre_grasp_pose_stamped = get_posestamped_msg(pre_position, orientation, camera_frame)                    
                grasp_pose_stamped = get_posestamped_msg(position, orientation, camera_frame)
                pour_pose_stamped = get_posestamped_msg(pour_position, pour_orientation, camera_frame)   
                        
                if vis:
                    mesh_resource = 'file://{}/common/assets/jaco_hand_mesh.ply'.format(src_path)
                    marker = get_marker_msg(position, orientation, 1.0, 1, mesh_resource, camera_frame, color=[0,1,0,1])
                    gripper_grasp_pose_topic.publish(marker)
                    marker = get_marker_msg(pour_position, pour_orientation, 1.0, 1, mesh_resource, camera_frame, color=[0,1,0,1])
                    gripper_pour_pose_topic.publish(marker)

                # grasp #
                if grasp_flag:

                    pose_list = [pre_grasp_pose_stamped, grasp_pose_stamped, pour_pose_stamped]
                    transformed_pose_list = transform_poselist_to_tgt(pose_list, robot_prefix+'link_base', tfBuffer)
                    pre_grasp_pose_stamped, grasp_pose_stamped, pour_pose_stamped = transformed_pose_list

                    # input of official function
                    pre_pose = get_pose_list_from_pose_stamped(pre_grasp_pose_stamped)
                    pose = get_pose_list_from_pose_stamped(grasp_pose_stamped)
                    pour_pose = get_pose_list_from_pose_stamped(pour_pose_stamped)
                    
                    open_size = {'mug':[50,50,50], 'bottle':[10,10,10], 'bowl':[50,50,50]}
                    if grasp_cfg['mug']['type'] == 'body':
                        open_size['mug'] = [10,10,10]
                    open_finger_positions = open_size[grasp_obj_cate]
                    close_finger_positions = [100,100,100]
           
                    try:
                        # open gripper
                        result = finger_action(open_finger_positions, robot_prefix)
                        # move to grasp pose
                        result = arm_action(pre_pose, robot_prefix)
                        # result = moveit_arm_plan(arm_group, pre_grasp_pose_stamped)
                        result = arm_action(pose, robot_prefix)
                        # close gripper to grasp object
                        result = finger_action(close_finger_positions, robot_prefix)
                        # take up object
                        result = arm_take_up_action(0.13, robot_prefix)
                        # move to bowl
                        result = arm_action(pour_pose, robot_prefix)
                        rospy.loginfo('results {}'.format(result))
                        # pour
                        pour_angle = 120 # in degree
                        if not right:
                            pour_angle = -pour_angle
                        result = arm_change_joint_action([0,0,0,0,0,pour_angle,0], robot_prefix)
                        result = arm_change_joint_action([0,0,0,0,0,-pour_angle,0], robot_prefix)
                        # move to place
                        # result = arm_take_up_action(0.05, robot_prefix)
                        result = arm_action(pose, robot_prefix)
                        # place
                        result = finger_action(open_finger_positions, robot_prefix)
                        result = arm_action(pre_pose, robot_prefix)
                        result = arm_take_up_action(0.20, robot_prefix)
                        # homeRobot(robot_prefix)
                        initial_joint_pose(robot_prefix)

                    except rospy.ROSInterruptException:
                        print('program interrupted before completion')
            
            else:
                rospy.loginfo('WARNING! No Grasp estimated!')
                continue
        else:
            rospy.loginfo('WARNING! No Object detected! Please check again!')
            continue

        # Grasp Only Once
        if grasp_flag:
            break
